package lk;


import java.util.HashMap;
import java.util.Map;

public class D230929 {
    static class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;
        TreeNode() {}
        TreeNode(int val) { this.val = val; }
        TreeNode(int val, TreeNode left, TreeNode right) {
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }

    // 统计二叉树中好节点的数目
    // 深度优先遍历
    // 时间：O(n)
    // 空间：O(n)
    public int goodNodes(TreeNode root) {
        int res = func(root, Integer.MIN_VALUE);
        return res;
    }

    public int func(TreeNode root, int max) {
        if (root == null) {
            return 0;
        }
        int count = 0;
        if (root.val >= max) {
            max = root.val;
            count++;
        }
        count += func(root.left, max);
        count += func(root.right, max);
        return count;
    }



    // 路径总和 III

    // 1.
    // 深度优先搜索
    // 穷举, 检测每一个节点作为起始节点, 向下延伸的所有路径
    // 时间：O(N^2)
    // 空间：O(N)
    public int pathSum(TreeNode root, int targetSum) {
        if (root == null) {
            return 0;
        }

        // 当前节点作为起始节点, 检测向下延伸的所有路径
        int ret = rootSum(root, targetSum);

        // 递归左右子树的节点
        ret += pathSum(root.left, targetSum);
        ret += pathSum(root.right, targetSum);

        return ret;
    }

    public int rootSum(TreeNode root, int targetSum) {
        int ret = 0;
        if (root == null) {
            return 0;
        }
        int val = root.val;
        if (val == targetSum) {
            ret++;
        }
        ret += rootSum(root.left, targetSum - val);
        ret += rootSum(root.right, targetSum - val);
        return ret;
    }


    // 2. 前缀和
    // 通过先序遍历二叉树, 记录下根节点 root 到当前节点 p 的路径上除当前节点以外所有节点的前缀和
    // 在已保存的路径前缀和 中查找是否存在前缀和刚好 等于 当前节点到根节点的前缀和 curr 减去 targetSum
    // 时间：O(N)
    // 空间：O(N)
    public int pathSum2(TreeNode root, int targetSum) {
        Map<Long, Integer> prefix = new HashMap<Long, Integer>();
        prefix.put(0L, 1);
        return dfs(root, prefix, 0, targetSum);
    }

    public int dfs(TreeNode root, Map<Long, Integer> prefix, long curr, int targetSum) {
        if (root == null) {
            return 0;
        }

        int ret = 0;
        curr += root.val;   // 当前节点的前缀和
        // prefix 记录了从根节点到当前节点, 路径上除了当前节点所有节点的前缀和

        ret = prefix.getOrDefault(curr - targetSum, 0);

        // 如果之前的前缀和中存在 curr - targetSum,
        // 则说明减去某个前缀和, 就存在一条路径和等于 targetSUm

        prefix.put(curr, prefix.getOrDefault(curr, 0) + 1); // 将当前节点的前缀和存入 prefix

        // 继续搜索下方的节点
        ret += dfs(root.left, prefix, curr, targetSum);
        ret += dfs(root.right, prefix, curr, targetSum);

        // 当前节点 以及 左右子树 都遍历完成了, 需要移除
        prefix.put(curr, prefix.getOrDefault(curr, 0) - 1);

        return ret;
    }
}
